#include <DCCpp.h>
 

// DCC 信号输出引脚配置
const int DCC_BRAKE_PIN = 9;    // DCC信号输出引脚
const int DCC_DIR_PIN = 12 ;    // 未使用
const int DCC_PWM_PIN = 3 ;     // 未使用
const int DCC_RX_PIN = A0;      // 轨道电流检测引脚

// 编程轨 信号输出引脚配置
const int PRO_BRAKE_PIN = 8;    // DCC信号输出引脚
const int PRO_DIR_PIN = 13 ;    // 未使用
const int PRO_PWM_PIN = 11 ;     // 未使用
const int PRO_RX_PIN = A1;      // 轨道电流检测引脚

// 当前轨道模式 (0=主轨道, 1=编程轨道)
int currentTrack = 0;
// 机车状态结构
struct LocoState {
  int addr=3;              // 地址
  uint8_t speed = 0;       // 当前速度 (0-126)
  bool direction = true;   // 方向 (true=前进, false=后退)
  FunctionsState functions; // 功能状态
};

// 当前机车的状态
LocoState locoState;

const int defaultAddress=3;

void setup() {
  // 初始化串口通信
  Serial.begin(115200);
  while (!Serial); // 等待串口连接
  
  Serial.println(F("DCC++ Controller v2.0 (Program Track Support)"));
  Serial.println(F("============================================"));
  Serial.println(F("Commands:"));
  Serial.println(F("  S[addr],[speed]   - Set speed (0-126)"));
  Serial.println(F("  D[addr],[dir]     - Set direction (0=Rev, 1=Fwd)"));
  Serial.println(F("  F[addr],[func],[state] - Set function (0-28)"));
  Serial.println(F("  E[addr]           - Emergency stop"));
  Serial.println(F("  R                 - Reset controller"));
  Serial.println(F("  I                 - Display system info"));
  Serial.println(F("  PON/POFF          - Turn track power on/off"));
  Serial.println(F("  T[track]          - Switch track (0=Main, 1=Prog)"));
  Serial.println(F("  RC[cv]            - Read CV value"));
  Serial.println(F("  WC[cv],[value]    - Write CV value"));
  Serial.println();
  // 初始化DCC++引擎
  DCCpp::begin();
  
  // 配置主轨道接口 - 使用头文件中定义的beginMain
  DCCpp::beginMain(DCC_DIR_PIN, DCC_BRAKE_PIN, DCC_PWM_PIN, DCC_RX_PIN);
  // 配置编程轨道接口
  DCCpp::beginProg(PRO_DIR_PIN, PRO_BRAKE_PIN, PRO_PWM_PIN, PRO_RX_PIN);
  
  // 启动DCC信号（主轨道）
  DCCpp::powerOn(true, false);
  
  // 显示初始状态
  displayStatus();
}

void loop() {
  // 必须持续调用以维持DCC信号
  DCCpp::loop();
  
  // 处理串口指令
  if (Serial.available() > 0) {
    processSerialCommand();
  }
  
  // 定期显示状态
  static unsigned long lastDisplayTime = 0;
  if (millis() - lastDisplayTime > 5000) {
    displayStatus();
    lastDisplayTime = millis();
  }
}

void processSerialCommand() {
  String command = Serial.readStringUntil('\n');
  command.trim();
  
  if (command.length() < 1) {
    Serial.println(F("ERR: Empty command"));
    return;
  }
  
  char cmdType = command.charAt(0);
  int addr = defaultAddress;  
  int value1 = 0;
  int value2 = 0;
  
  // 尝试解析地址
  if (command.length() > 1 && isDigit(command.charAt(1))) {
    int firstComma = command.indexOf(',');
    if (firstComma > 1) {
      addr = command.substring(1, firstComma).toInt();
    } else {
      addr = command.substring(1).toInt();
    }
  }
  
  // 解析命令参数
  int firstComma = command.indexOf(',');
  int secondComma = (firstComma != -1) ? command.indexOf(',', firstComma + 1) : -1;
  
  if (cmdType == 'S' || cmdType == 'D' || cmdType == 'E' || 
      cmdType == 'T' || cmdType == 'R' || cmdType == 'W') {
    if (firstComma != -1) {
      value1 = command.substring(firstComma + 1).toInt();
    } else if (cmdType == 'T' || cmdType == 'R') {
      value1 = command.substring(1).toInt();
    }
  } 
  else if (cmdType == 'F' || cmdType == 'W') {
    if (firstComma != -1) {
      value1 = command.substring(firstComma + 1, secondComma).toInt();
      if (secondComma != -1) {
        value2 = command.substring(secondComma + 1).toInt();
      }
    }
  }
  
  // 执行命令
  switch (cmdType) {
    case 'S': // 速度命令
      if (addr < 1 || addr > 127) {
        Serial.println(F("ERR: Invalid address (1-127)"));
        return;
      }
      if (value1 < 0 || value1 > 126) {
        Serial.println(F("ERR: Speed out of range (0-126)"));
        return;
      } 
      setLocoSpeed(addr, value1);
      break;
      
    case 'D': // 方向命令
      if (addr < 1 || addr > 127) {
        Serial.println(F("ERR: Invalid address (1-127)"));
        return;
      }
      if (value1 != 0 && value1 != 1) {
        Serial.println(F("ERR: Direction must be 0 or 1"));
        return;
      } 
      setLocoDirection(addr, value1);
      break;
      
    case 'F': // 功能命令
      if (addr < 1 || addr > 127) {
        Serial.println(F("ERR: Invalid address (1-127)"));
        return;
      }
      if (value1 < 0 || value1 > 28) {
        Serial.println(F("ERR: Function out of range (0-28)"));
        return;
      }
      if (value2 != 0 && value2 != 1) {
        Serial.println(F("ERR: State must be 0 or 1"));
        return;
      } 
      setLocoFunction(addr, value1, value2);
      break;
      
    case 'E': // 紧急停止
      if (addr < 1 || addr > 127) {
        Serial.println(F("ERR: Invalid address (1-127)"));
        return;
      }
      emergencyStop(addr);
      break;
      
    case 'R': // 读取CV值
      if (command.charAt(1) == 'C') {
        // RC命令：读取CV值
        readCV(value1);
      } else {
        // R命令：复位控制器
        resetController();
      }
      break;
      
    case 'W': // 写入CV值
      if (command.charAt(1) == 'C') {
        // WC命令：写入CV值
        writeCV(value1, value2);
      }
      break;
      
    case 'T': // 切换轨道
      if (value1 == 0 || value1 == 1) {
        switchTrack(value1);
      } else {
        Serial.println(F("ERR: Invalid track (0=Main, 1=Prog)"));
      }
      break;
      
    case 'I': // 显示系统信息
      displaySystemInfo();
      break;
      
    case 'P': // 电源控制
      if (command.substring(1) == "ON") {
        DCCpp::powerOn(true, true);
        Serial.println(F("Track power: ON (both tracks)"));
      } else if (command.substring(1) == "OFF") {
        DCCpp::powerOff(true, true);
        Serial.println(F("Track power: OFF (both tracks)"));
      } else if (command.substring(1) == "MAIN") {
        DCCpp::powerOn(true, false);
        Serial.println(F("Main track power: ON"));
      } else if (command.substring(1) == "PROG") {
        DCCpp::powerOn(false, true);
        Serial.println(F("Prog track power: ON"));
      }
      break;
      
    default:
      Serial.println(F("ERR: Unknown command"));
      Serial.println(F("Valid commands: S, D, F, E, R, I, P, T, RC, WC"));
      break;
  }
}

// 设置机车速度
void setLocoSpeed(int addr, int speed) {
  // 更新本地状态
  locoState.addr=addr;
  locoState.speed = speed;
  
  // 根据当前轨道选择设置函数
  if (currentTrack == 0) {
    // 主轨道
    DCCpp::setSpeedMain(1, locoState.addr, 128, locoState.speed , locoState.direction);
  } else {
    // 编程轨道
    DCCpp::setSpeedProg(1, locoState.addr, 128, locoState.speed , locoState.direction);
  }
  
  // 显示状态
  Serial.print(F("Loco "));
  Serial.print(locoState.addr);
  Serial.print(F(" on "));
  Serial.print(currentTrack == 0 ? "Main" : "Prog");
  Serial.print(F(" track: Speed="));
  Serial.print(locoState.speed);
  Serial.print(F(", Dir="));
  Serial.println(locoState.direction ? F("FWD") : F("REV"));
}

// 设置机车方向
void setLocoDirection(int addr, int direction) {
  // 更新本地状态
  locoState.addr=addr;
  locoState.direction = direction;
  
  // 根据当前轨道选择设置函数
  if (currentTrack == 0) {
    // 主轨道
    DCCpp::setSpeedMain(1, locoState.addr, 128, locoState.speed, locoState.direction );
  } else {
    // 编程轨道
    DCCpp::setSpeedProg(1, locoState.addr, 128, locoState.speed, locoState.direction );
  }
  
  // 显示状态
  Serial.print(F("Loco "));
  Serial.print(locoState.addr);
  Serial.print(F(" on "));
  Serial.print(currentTrack == 0 ? "Main" : "Prog");
  Serial.print(F(" track: Dir="));
  Serial.print(locoState.direction  ? F("FWD") : F("REV"));
  Serial.print(F(", Speed="));
  Serial.println(locoState.speed);
}

// 设置机车功能
void setLocoFunction(int addr, int function, int state) {
  // 更新功能状态
  locoState.addr=addr;
  if (state) {
    locoState.functions.activate(function);
  } else {
    locoState.functions.inactivate(function);
  }
  
  // 根据当前轨道选择设置函数
  if (currentTrack == 0) {
    // 主轨道
    DCCpp::setFunctionsMain(1, locoState.addr, locoState.functions);
  } else {
    // 编程轨道
    DCCpp::setFunctionsProg(1, locoState.addr, locoState.functions);
  }
  
  // 显示状态
  Serial.print(F("Loco "));
  Serial.print(locoState.addr);
  Serial.print(F(" on "));
  Serial.print(currentTrack == 0 ? "Main" : "Prog");
  Serial.print(F(" track: F"));
  Serial.print(function);
  Serial.print(F("="));
  Serial.println(state ? F("ON") : F("OFF"));
}

// 紧急停止
void emergencyStop(int addr) {
  locoState.addr=addr;
  // 根据当前轨道选择设置函数
  if (currentTrack == 0) {
    // 主轨道
    DCCpp::setSpeedMain(1, locoState.addr, 128, 1, locoState.direction);
  } else {
    // 编程轨道
    DCCpp::setSpeedProg(1, locoState.addr, 128, 1, locoState.direction);
  }
  
  // 更新本地状态
  locoState.speed = 0;
  
  Serial.print(F("EMERGENCY STOP for loco "));
  Serial.print(locoState.addr);
  Serial.print(F(" on "));
  Serial.println(currentTrack == 0 ? "Main" : "Prog");
}

// 切换轨道
void switchTrack(int track) {
  if (track == currentTrack) {
    Serial.print(F("Already on "));
    Serial.println(track == 0 ? "Main track" : "Prog track");
    return;
  }
  
  // 关闭当前轨道电源
  if (currentTrack == 0) {
    DCCpp::powerOff(true, false);
  } else {
    DCCpp::powerOff(false, true);
  }
  
  // 切换轨道
  currentTrack = track;
  
  // 开启新轨道电源
  if (currentTrack == 0) {
    DCCpp::powerOn(true, false);
    Serial.println(F("Switched to Main track"));
  } else {
    DCCpp::powerOn(false, true);
    Serial.println(F("Switched to Prog track"));
  }
}

// 读取CV值
void readCV(int cv) {
  if (cv < 1 || cv > 1024) {
    Serial.println(F("ERR: CV out of range (1-1024)"));
    return;
  }
  
  int value = -1;
  
  if (currentTrack == 0) {
    // 主轨道读取
    value = DCCpp::readCvMain(cv);
  } else {
    // 编程轨道读取
    value = DCCpp::readCvProg(cv);
  }
  
  if (value >= 0) {
    Serial.print(F("CV"));
    Serial.print(cv);
    Serial.print(F(" value: "));
    Serial.println(value);
  } else {
    Serial.print(F("Failed to read CV"));
    Serial.println(cv);
  }
}

// 写入CV值
void writeCV(int cv, int value) {
  if (cv < 1 || cv > 1024) {
    Serial.println(F("ERR: CV out of range (1-1024)"));
    return;
  }
  
  if (value < 0 || value > 255) {
    Serial.println(F("ERR: CV value out of range (0-255)"));
    return;
  }
  
  bool success = false;
  
  if (currentTrack == 0) {
    // 主轨道写入
    DCCpp::writeCvMain(cv, (byte)value);
    success = true; // 主轨道写入没有返回状态
  } else {
    // 编程轨道写入
    success = DCCpp::writeCvProg(cv, (byte)value);
  }
  
  if (success) {
    Serial.print(F("CV"));
    Serial.print(cv);
    Serial.print(F(" set to "));
    Serial.println(value);
  } else {
    Serial.print(F("Failed to write CV"));
    Serial.println(cv);
  }
}

// 复位控制器
void resetController() {
  // 停止当前机车
  if (currentTrack == 0) {
    DCCpp::setSpeedMain(1, locoState.addr, 128, 0, true);
  } else {
    DCCpp::setSpeedProg(1, locoState.addr, 128, 0, true);
  }
  
  // 重置状态
  locoState.speed = 0;
  locoState.direction = true;
  locoState.functions.clear();
  locoState.addr=defaultAddress;

  // 重置轨道
  currentTrack = 0;
  DCCpp::powerOn(true, false);
  
  Serial.println(F("Controller reset complete"));
}

// 显示系统信息
void displaySystemInfo() {
  Serial.println(F("\n===== SYSTEM INFO ====="));
  Serial.print(F("Main Track Signal Pin: "));
  Serial.println(DCC_BRAKE_PIN);
  Serial.print(F("Main Track Current Pin: "));
  Serial.println(DCC_RX_PIN);
  Serial.print(F("Prog Track Signal Pin: "));
  Serial.println(PRO_BRAKE_PIN);
  Serial.print(F("Prog Track Current Pin: "));
  Serial.println(PRO_RX_PIN);
  
  Serial.print(F("Main Track Power: "));
  Serial.println(DCCpp::IsPowerOnMain ? F("ON") : F("OFF"));
  Serial.print(F("Prog Track Power: "));
  Serial.println(DCCpp::IsPowerOnProg ? F("ON") : F("OFF"));
  
  Serial.print(F("Main Track Current: "));
  Serial.print(DCCpp::getCurrentMain());
  Serial.println(F(" mA"));
  
  Serial.print(F("Prog Track Current: "));
  Serial.print(DCCpp::getCurrentProg());
  Serial.println(F(" mA"));
  
  Serial.print(F("Current Track: "));
  Serial.println(currentTrack == 0 ? "Main" : "Prog");
  
  Serial.println(F("======================\n"));
}

// 显示当前状态
void displayStatus() {
  Serial.println(F("\n===== STATUS ====="));
  Serial.print(F("Current loco: "));
  Serial.println(locoState.addr);
  
  Serial.print(F("Current track: "));
  Serial.println(currentTrack == 0 ? "Main" : "Prog");
  
  Serial.print(F("Speed: "));
  Serial.print(locoState.speed);
  Serial.print(F("/126, Direction: "));
  Serial.println(locoState.direction ? F("Forward") : F("Reverse"));
  
  Serial.print(F("Main Track Power: "));
  Serial.println(DCCpp::IsPowerOnMain ? F("ON") : F("OFF"));
  Serial.print(F("Main Track Current: "));
  Serial.print(DCCpp::getCurrentMain());
  Serial.println(F(" mA"));
  
  Serial.print(F("Prog Track Power: "));
  Serial.println(DCCpp::IsPowerOnProg ? F("ON") : F("OFF"));
  Serial.print(F("Prog Track Current: "));
  Serial.print(DCCpp::getCurrentProg());
  Serial.println(F(" mA"));
  
  Serial.println(F("==================\n"));
}